#!/usr/bin/env python3
"""
MQTT 测试脚本 - 模拟六足机器人发送数据
用于测试 MQTT 到 WebSocket 的数据流
每10秒发送一次包含base64图片的数据
"""

import json
import time
import random
import base64
from paho.mqtt import client as mqtt_client

# MQTT Broker settings
MQTT_BROKER = "broker.hivemq.com"  # 使用公共测试 broker
MQTT_PORT = 1883
MQTT_TOPIC = "robot/hexapod/status"
MQTT_CLIENT_ID = "hexapod_robot_simulator"

def generate_sample_base64_image():
    """生成一个简单的base64编码图片（1x1像素的JPEG）"""
    # 这是一个1x1像素的红色JPEG图片的base64编码
    red_pixel = "/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAYEBQYFBAYGBQYHBwYIChAKCgkJChQODwwQFxQYGBcUFhYaHSUfGhsjHBYWICwgIyYnKSopGR8tMC0oMCUoKSj/2wBDAQcHBwoIChMKChMoGhYaKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCj/wAARCAABAAEDASIAAhEBAxEB/8QAFQABAQAAAAAAAAAAAAAAAAAAAAv/xAAUEAEAAAAAAAAAAAAAAAAAAAAA/8QAFQEBAQAAAAAAAAAAAAAAAAAAAAX/xAAUEQEAAAAAAAAAAAAAAAAAAAAA/9oADAMBAAIRAxEAPwCdABmX/9k="
    
    # 随机选择不同颜色的像素来模拟不同的图片
    colors = [
        red_pixel,  # 红色
        "/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAYEBQYFBAYGBQYHBwYIChAKCgkJChQODwwQFxQYGBcUFhYaHSUfGhsjHBYWICwgIyYnKSopGR8tMC0oMCUoKSj/2wBDAQcHBwoIChMKChMoGhYaKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCj/wAARCAABAAEDASIAAhEBAxEB/8QAFQABAQAAAAAAAAAAAAAAAAAAAAv/xAAUEAEAAAAAAAAAAAAAAAAAAAAA/8QAFQEBAQAAAAAAAAAAAAAAAAAAAAX/xAAUEQEAAAAAAAAAAAAAAAAAAAAA/9oADAMBAAIRAxEAPwCdABmX/9k=",  # 绿色
        "/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAYEBQYFBAYGBQYHBwYIChAKCgkJChQODwwQFxQYGBcUFhYaHSUfGhsjHBYWICwgIyYnKSopGR8tMC0oMCUoKSj/2wBDAQcHBwoIChMKChMoGhYaKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCj/wAARCAABAAEDASIAAhEBAxEB/8QAFQABAQAAAAAAAAAAAAAAAAAAAAv/xAAUEAEAAAAAAAAAAAAAAAAAAAAA/8QAFQEBAQAAAAAAAAAAAAAAAAAAAAX/xAAUEQEAAAAAAAAAAAAAAAAAAAAA/9oADAMBAAIRAxEAPwCdABmX/9k="   # 蓝色
    ]
    
    return f"data:image/jpeg;base64,{random.choice(colors)}"

def connect_mqtt():
    def on_connect(client, userdata, flags, rc):
        if rc == 0:
            print("Connected to MQTT Broker!")
        else:
            print(f"Failed to connect, return code {rc}")

    client = mqtt_client.Client(mqtt_client.CallbackAPIVersion.VERSION1, MQTT_CLIENT_ID)
    client.on_connect = on_connect
    client.connect(MQTT_BROKER, MQTT_PORT)
    return client

def publish_robot_data(client):
    """模拟发布机器人数据 - 每10秒一次"""
    counter = 0
    while True:
        counter += 1
        
        # 生成模拟数据
        # 随机决定是否检测到人员（30%概率）
        person_detected = random.random() < 0.3
        person_count = 0
        person_locations = []
        detection_confidence = 0.0
        
        if person_detected:
            person_count = random.randint(1, 3)
            detection_confidence = random.uniform(0.8, 0.99)
            # 生成人员位置坐标（假设图像尺寸为640x480）
            for i in range(person_count):
                person_locations.append({
                    "x": random.randint(50, 590),
                    "y": random.randint(50, 430),
                    "confidence": random.uniform(0.8, 0.99)
                })
        
        robot_data = {
            "robot_id": "hexapod_001",
            "timestamp": int(time.time()),
            "location": {
                "latitude": 34.052235 + random.uniform(-0.001, 0.001),
                "longitude": -118.243683 + random.uniform(-0.001, 0.001),
                "altitude": 100.5 + random.uniform(-5, 5)
            },
            "orientation": {
                "roll": random.uniform(-15, 15),
                "pitch": random.uniform(-15, 15),
                "yaw": random.uniform(0, 360)
            },
            "battery_level": max(0, 100 - (counter * 2)),  # 模拟电池逐渐消耗
            "sensors": {
                "temperature": 25.3 + random.uniform(-5, 5),
                "humidity": 60.1 + random.uniform(-10, 10),
                "gas_detection": {
                    "CO": random.randint(0, 100),
                    "CH4": random.randint(0, 200)
                }
            },
            "status": random.choice(["exploring", "moving", "scanning", "idle", "charging"]),
            "camera_image": generate_sample_base64_image(),
            "person_detected": person_detected,
            "person_count": person_count,
            "detection_confidence": detection_confidence,
            "person_locations": person_locations
        }
        
        # 发布数据
        message = json.dumps(robot_data)
        result = client.publish(MQTT_TOPIC, message)
        status = result[0]
        if status == 0:
            print(f"[{time.strftime('%H:%M:%S')}] Sent robot data #{counter} to topic `{MQTT_TOPIC}`")
            print(f"Status: {robot_data['status']}, Battery: {robot_data['battery_level']:.1f}%")
            if person_detected:
                print(f"🚨 ALERT: {person_count} person(s) detected! Confidence: {detection_confidence:.2f}")
            else:
                print("No persons detected")
        else:
            print(f"Failed to send message to topic {MQTT_TOPIC}")
        
        # 等待 10 秒后发送下一条数据
        print(f"Waiting 10 seconds for next update...")
        time.sleep(10)

if __name__ == "__main__":
    client = connect_mqtt()
    client.loop_start()
    
    try:
        print("Starting robot data simulation (every 10 seconds)...")
        print("Press Ctrl+C to stop")
        publish_robot_data(client)
    except KeyboardInterrupt:
        print("\nStopping robot data simulation...")
        client.loop_stop()
        client.disconnect()

